import os
import torch
import torch.distributed as dist
from diffusers import StableVideoDiffusionPipeline
from PIL import Image
import requests
from io import BytesIO
import psutil
from moviepy import ImageSequenceClip

def print_resource_usage():
    # GPU内存使用
    if torch.cuda.is_available():
        print(f"GPU Memory - Allocated: {torch.cuda.memory_allocated()/1024**2:.2f}MB, "
              f"Cached: {torch.cuda.memory_reserved()/1024**2:.2f}MB")
    # CPU使用率
    print(f"CPU Usage: {psutil.cpu_percent()}%")

def setup(rank, world_size, master_ip):
    os.environ['MASTER_ADDR'] = master_ip  # 主节点地址
    os.environ['MASTER_PORT'] = '8989'  # 主节点端口

    # 初始化进程组，将后端改为gloo
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def main(rank, world_size, master_ip):
    setup(rank, world_size, master_ip)

    # 设置代理（如果需要）
    # os.environ['HTTP_PROXY'] = 'http://your_proxy_address:your_proxy_port'
    # os.environ['HTTPS_PROXY'] = 'http://your_proxy_address:your_proxy_port'

    # 加载初始图像
    url = "https://pic.rmb.bdstatic.com/bjh/other/7132183c3a3973c03796a833e857df67.jpeg?for=bg"  # 替换为实际图像URL
    response = requests.get(url)
    init_image = Image.open(BytesIO(response.content)).convert("RGB")
    if rank == 0:
        print('loading image...')
        print("Initial resource usage:")
        print_resource_usage()

    pipeline = StableVideoDiffusionPipeline.from_pretrained(
        "stabilityai/stable-video-diffusion-img2vid-xt",  # 若使用本地模型，修改为本地路径
        device=f"cuda:{rank}",
        torch_dtype=torch.float32,  # 降低精度要求
        cache_dir='models')

    if rank == 0:
        print("\nAfter model loading:")
        print_resource_usage()

    if rank == 0:
        print('generating video...')
    try:
        if rank == 0:
            print("\nBefore video generation:")
            print_resource_usage()
        video_frames = pipeline(init_image, num_frames=8).frames
        # 清理GPU缓存
        torch.cuda.empty_cache()
        if rank == 0:
            print("\nAfter video generation:")
            print_resource_usage()
    finally:
        if rank == 0:
            print('saving video...')
        # 使用moviepy将帧保存为视频
        if rank == 0:
            clip = ImageSequenceClip([frame for frame in video_frames], fps=4)  # 设置帧率为4fps
            clip.write_videofile("videos/output.mp4", codec="libx264")

    cleanup()

if __name__ == "__main__":
    world_size = 2  # 机器数量
    master_ip = "192.168.12.251"  # 主节点的 IP 地址，需要根据实际情况修改
    torch.multiprocessing.spawn(main,
                                args=(world_size, master_ip),
                                nprocs=1,
                                join=True)